from absl import app, flags
import h5py
import os
import pandas
import socket

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import coin.strategy.mm.tool.archive_base as abase
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2


_fi_repo = '/remote/iosg/data-2/buckets/feed.derived.interval_h5'
_fi_root = f'{_fi_repo}/coin/main'
_fi_resol = 'PT1M'
fi_dir = f'{_fi_root}/{_fi_resol}'


def load_mea(mea, trading_date):
  tdstr = trading_date.strftime("%Y%m%d")
  h5data = h5py.File(f'{fi_dir}/{mea}/{tdstr}/{mea}--ohlc.h5', 'r')
  nativesymbolstrs = h5data['universe'][:]
  mt, ex, api = mea.split(".")
  iidx = []
  for nativesymbolstr in nativesymbolstrs:
    try:
      product = generate_product_from_str2(mt, ex, api, nativesymbolstr, trading_date)
      iidx.append(
          product.subscription_symbol.replace("CURRENT_", "").replace("THIS_", "")
          if hasattr(product, 'subscription_symbol')
          else product.symbol)
    except Exception:
      iidx.append(nativesymbolstr)
  tidx = pandas.DatetimeIndex(h5data['timestamp'][:])
  dfdict = {}
  for key in h5data.keys():
    if key not in ['timestamp', 'universe']:
      df = pandas.DataFrame(h5data[key][:])
      df.index = tidx
      df.columns = iidx
      dfdict[key] = df
  return dfdict


def setup_plt():
  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['figure.figsize'] = 8, 4
  plt.rcParams['font.size'] = 6
  plt.rcParams['legend.fontsize'] = 6
  plt.rcParams['xtick.labelsize'] = 6
  plt.rcParams['ytick.labelsize'] = 6


def get_ip_address():
  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  s.connect(("8.8.8.8", 80))
  return s.getsockname()[0]


ip_addr = get_ip_address()
cwd = os.getcwd()


def plot_and_save(seri1, seri2, legends, filename):
  plt.plot(seri1, 'r-', seri2, 'g-', lw=0.5, drawstyle='steps-post')
  plt.legend(legends)
  plt.twinx()
  plt.plot(seri1 - seri2, 'b-', lw=0.5, drawstyle='steps-post', alpha=0.5)
  plt.xticks(rotation=60)
  dest = os.path.join(flags.FLAGS.out_dir, filename)
  print(f"http://{ip_addr}:8080" + os.path.join(cwd, dest))
  plt.savefig(dest, dpi=200)
  plt.close()


def main(_):
  trading_dates = abase.get_trading_dates(flags.FLAGS.trading_date)
  os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
  mea1, prod1 = flags.FLAGS.prod1.split(":")
  mea2, prod2 = flags.FLAGS.prod2.split(":")
  seri1s = []
  seri2s = []
  for trading_date in trading_dates:
    dfs1 = load_mea(mea1, trading_date)
    dfs2 = load_mea(mea2, trading_date)
    seri1 = dfs1['OPEN_MID'][prod1]
    seri2 = dfs2['OPEN_MID'][prod2]
    seri1s.append(seri1)
    seri2s.append(seri2)
    if flags.FLAGS.mode == 'plot':
      setup_plt()
      plot_and_save(
          seri1,
          seri2,
          [f'{mea1}:{prod1}', f'{mea2}:{prod2}'],
          f"{trading_date.strftime('%Y%m%d')}_{mea1}_{mea2}.png")
  wseri1 = pandas.concat(seri1s, axis=0)
  wseri2 = pandas.concat(seri2s, axis=0)
  if flags.FLAGS.mode == 'plot':
    plot_and_save(
        wseri1,
        wseri2,
        [f'{mea1}:{prod1}', f'{mea2}:{prod2}'],
        f"{flags.FLAGS.trading_date}_{mea1}_{mea2}.png")
  elif flags.FLAGS.mode == 'pdb':
    import pdb
    pdb.set_trace()


# ./pyrunner python/coin/research/feed_interval_plot.py --trading_date=20201126
if __name__ == "__main__":
  flags.DEFINE_string('out_dir', 'pic_fip', '')
  flags.DEFINE_string('prod1', 'Futures.Huobi.v1:BTC-USD.QUARTER', '')
  flags.DEFINE_string('prod2', 'Futures.Okex.v3:BTC-USD.QUARTER', '')
  flags.DEFINE_string('mode', 'plot', 'plot,pdb')
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  app.run(main)
